package com.xiashitech.agent.instrumentation.internal.classloader.inst;

import com.xiashitech.agent.boot.LoadClassHelper;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.BooleanMatcher;
import net.bytebuddy.matcher.ElementMatcher;

import java.io.File;
import java.io.InputStream;
import java.lang.reflect.Method;
import java.security.CodeSource;
import java.security.SecureClassLoader;
import java.security.cert.Certificate;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;

import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.extendsClass;
import static net.bytebuddy.matcher.ElementMatchers.*;

/**
 * 加载非隔离的，运行在业务系统内部的agent帮助类
 */
public class XSLoadInjectedClassInstrumentation implements TypeInstrumentation {

  @Override
  public ElementMatcher<ClassLoader> classLoaderOptimization() {
    try{
      Class.forName("jdk.internal.loader.ClassLoaders"); //仅在jdk9+情况下启用
      return BooleanMatcher.of(true);
    }catch (Throwable t){
      return BooleanMatcher.of(false);
    }
  }

  @Override
  public ElementMatcher<TypeDescription> typeMatcher() {
    return extendsClass(named("java.lang.ClassLoader"));
  }

  @Override
  public void transform(TypeTransformer transformer) {
    transformer.applyAdviceToMethod(
        isMethod()
            .and(named("loadClass"))
            .and(
                takesArguments(1)
                    .and(takesArgument(0, String.class))
                    .or(
                        takesArguments(2)
                            .and(takesArgument(0, String.class))
                            .and(takesArgument(1, boolean.class))))
            .and(isPublic().or(isProtected()))
            .and(not(isStatic())),
        XSLoadInjectedClassInstrumentation.class.getName() + "$XSLoadClassAdvice");
  }

  @SuppressWarnings("unused")
  public static class XSLoadClassAdvice {

    @Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class)
    public static Class<?> onEnter(
        @Advice.This ClassLoader classLoader, @Advice.Argument(0) String name) {
      if(name != null && (name.startsWith("com.xiashitech.") && !name.contains(".boot."))
              || name.startsWith("xiashitech.shade.")) {
        try {
          ClassLoader targetLoader = classLoader;
          if(LoadClassHelper.toPlatformLoader(name)) {
            while (targetLoader.getParent() != null) {
              targetLoader = targetLoader.getParent();
            }
          }

          Method findLoadedClassMethod = ClassLoader.class.getDeclaredMethod("findLoadedClass", String.class);
          findLoadedClassMethod.setAccessible(true);
          synchronized (targetLoader) {
            Class xsCls = (Class) findLoadedClassMethod.invoke(targetLoader, name);
            if (xsCls != null)
              return xsCls;

            JarFile jarFile = LoadClassHelper.getJarFile();
            JarEntry jarEntry = jarFile.getJarEntry(name.replace('.', '/') + ".class");
            int size = (int) jarEntry.getSize();
            byte[] bytes = new byte[size];
            try (InputStream is = jarFile.getInputStream(jarEntry)) {
              int offset = 0;
              int read;
              while (offset < size && (read = is.read(bytes, offset, size - offset)) != -1) {
                offset += read;
              }
            }
            CodeSource codeSource = new CodeSource(new File(LoadClassHelper.getAgentPath()).toURI().toURL(), (Certificate[]) null);
            Method defineClassMethod = SecureClassLoader.class.getDeclaredMethod("defineClass", String.class, byte[].class, int.class, int.class, CodeSource.class);
            defineClassMethod.setAccessible(true);
            xsCls = (Class) defineClassMethod.invoke(targetLoader, name, bytes, 0, bytes.length, codeSource);
            return xsCls;
          }
        } catch (Throwable e) {

        }
      }

//      try {
//        if(classLoader instanceof SecureClassLoader && name != null && (name.startsWith("com.xiashitech") && name.contains("extload") && !name.contains("extloadproxy") || name.startsWith("xiashitech.shade.extload"))) {
//          Method findLoadedClassMethod = ClassLoader.class.getDeclaredMethod("findLoadedClass", String.class);
//          findLoadedClassMethod.setAccessible(true);
//          synchronized (classLoader) {
//            Class extloadClass = (Class) findLoadedClassMethod.invoke(classLoader, name);
//            if (extloadClass != null) {
//              return extloadClass;
//            }
//            JarFile jarFile = AgentCache.getJarFile();
//            JarEntry jarEntry = jarFile.getJarEntry(name.replace('.', '/') + ".classdata");
//            int size = (int) jarEntry.getSize();
//            byte[] bytes = new byte[size];
//            try (InputStream is = jarFile.getInputStream(jarEntry)) {
//              int offset = 0;
//              int read;
//              while (offset < size && (read = is.read(bytes, offset, size - offset)) != -1) {
//                offset += read;
//              }
//            }
//            CodeSource codeSource = new CodeSource(new File(AgentCache.getAgentPath()).toURI().toURL(), (Certificate[]) null);
//            Method defineClassMethod = SecureClassLoader.class.getDeclaredMethod("defineClass", String.class, byte[].class, int.class, int.class, CodeSource.class);
//            defineClassMethod.setAccessible(true);
//            extloadClass = (Class) defineClassMethod.invoke(classLoader, name, bytes, 0, bytes.length, codeSource);
//            return extloadClass;
//          }
//        }
//      } catch (Throwable e) {
//        StartConfig.throwableStackPrinter(e);
//      }

      return null;
    }

    @Advice.OnMethodExit(onThrowable = Throwable.class)
    public static void onExit(
        @Advice.Return(readOnly = false) Class<?> result, @Advice.Enter Class<?> loadedClass) {
      if (loadedClass != null) {
        result = loadedClass;
      }
    }
  }
}
